import time
import threading
from queue import Queue, Empty
from contextlib import contextmanager
from DrissionPage import Chromium, ChromiumOptions


class TabPool:
    """Tab池管理器，类似线程池的概念"""
    
    def __init__(self, max_tabs=5, browser_options=None):
        self.max_tabs = max_tabs
        self.browser = self._create_browser(browser_options)
        self.tab_queue = Queue(maxsize=max_tabs)
        self.lock = threading.Lock()
        self.closed = False
        
        # 初始化tab池
        self._init_tabs()
    
    def _create_browser(self, options=None):
        """创建浏览器实例"""
        if options is None:
            co = ChromiumOptions()
            co.set_timeouts(base=5, page_load=30)
            co.auto_port()
        else:
            co = options
        return Chromium(addr_or_opts=co)
    
    def _init_tabs(self):
        """初始化tab池"""
        for i in range(self.max_tabs):
            if i == 0:
                tab = self.browser.get_tab()
            else:
                tab = self.browser.new_tab()
            self.tab_queue.put(tab)
    
    def get_tab(self, timeout=None):
        """获取一个可用的tab"""
        if self.closed:
            raise RuntimeError("TabPool已关闭")
        
        try:
            tab = self.tab_queue.get(timeout=timeout)
            return tab
        except Empty:
            raise TimeoutError("获取tab超时")
    
    def return_tab(self, tab):
        """归还tab到池中"""
        if not self.closed:
            self.tab_queue.put(tab)
    
    @contextmanager
    def acquire_tab(self, timeout=None):
        """上下文管理器方式获取tab"""
        tab = self.get_tab(timeout)
        try:
            yield tab
        finally:
            self.return_tab(tab)
    
    def close(self):
        """关闭tab池和浏览器"""
        with self.lock:
            if not self.closed:
                self.closed = True
                self.browser.quit()
    
    def __enter__(self):
        return self
    
    def __exit__(self, exc_type, exc_val, exc_tb):
        self.close()

class TabPoolManager:
    """Tab池管理器的增强版本"""
    
    def __init__(self, max_tabs=5, browser_options=None):
        self.pool = TabPool(max_tabs, browser_options)
        self.active_tasks = 0
        self.lock = threading.Lock()
    
    def submit_task(self, func, *args, **kwargs):
        """提交任务到tab池"""
        def wrapper():
            with self.lock:
                self.active_tasks += 1
            
            try:
                with self.pool.acquire_tab() as tab:
                    return func(tab, *args, **kwargs)
            finally:
                with self.lock:
                    self.active_tasks -= 1
        
        thread = threading.Thread(target=wrapper)
        thread.start()
        return thread
    
    def wait_all_tasks(self):
        """等待所有任务完成"""
        while True:
            with self.lock:
                if self.active_tasks == 0:
                    break
            time.sleep(0.1)
    
    def close(self):
        """关闭管理器"""
        self.pool.close()